package datastructure.tree.avl2;

/**
 * AVL树
 *
 * 旋转的节点颜色说明:
 *      red: 需要向下的节点
 *      yellow: 需要向上的节点
 *      green: 需要换父节点的节点
 *
 * @author 杨 强
 * @package datastructure.tree.avl2
 * @createTime 2023-03-30 22:44
 */
public class AVL {
    public static void main(String[] args) {
        AVL avl = new AVL();
        avl.put(3, 3);
        avl.put(1, 1);
        avl.put(4, 4);
        avl.remove(1);
        System.out.println(avl);
    }

    /**
     * 跟节点
     */
    AVLNode root;


    /**
     * 删除节点
     * @param key 删除节点的key
     */
    public void remove(int key){
        root = doRemove(root, key);
    }

    /**
     * 删除节点
     * @param node 是否是删除的节点, 第一次循环是根节点
     * @param key 节点key
     * @return 新的根节点
     */
    private AVLNode doRemove(AVLNode node, int key){
        // 1 node == null
        if (node == null){
            return node;
        }

        // 2 没有找到key
        if (key < node.key){
            node.left = doRemove(node.left, key);
        }else if (key > node.key){
            node.right = doRemove(node.right, key);
        }else {
            // 3 找到key, 但是有3种情况: 1-没有孩子, 2-有一个孩子, 3-有2个孩子
            if (node.left == null && node.right == null){
                return null;
            }else if (node.left == null){
                node = node.right;
            }else if (node.right == null){
                node = node.left;
            }else {
                // 被删除节点有2个孩子, 是最复杂的情况, 被删除的节点有2个孩子, 看图写代码.
                // node是被删除节点
                // 先找到被删除节点的后继节点,往右子树找到头
                AVLNode s = node.right;
                while (s.left != null){
                    s = s.left;
                }
                // 处理后继节点, 也就是在被删除节点中的右子树中把后继节点删除掉
                // 删除节点返回的是这个子树的根节点, 也就是把被删除节点的右子树赋值给后继节点的右子树
                s.right = doRemove(node.right, s.key);
                // 后继节点的左子树就是被删除节点的左子树
                s.left = node.left;
                // 后继节点代替被删除节点
                node = s;
            }
        }

        // 4 更新高度
        updateHeight(node);

        // 5 检查是否平衡
        return balance(node);
    }


    /**
     * 添加元素, 看图一步一步写.
     * @param key 元素key
     * @param value 元素值
     */
    public void put(int key, Object value){
        root = doPut(root, key, value);
    }

    /**
     * 创建需要添加的节点,并且把添加的节点放到合适的位置
     * @param node 跟节点
     * @param key key
     * @param value 值
     * @return 带有新值的节点
     */
    private AVLNode doPut(AVLNode node, int key, Object value){
        // 1 找到空位, 创建新节点
        if (node == null){
            return new AVLNode(key, value);
        }

        // 2 key已存在, 更新
        if (node.key == key){
            node.value = value;
            return node;
        }

        // 3 继续查找
        if (key < node.key){
            node.left = doPut(node.left, key, value);
        }else {
            node.right = doPut(node.right, key, value);
        }

        // 更新节点高度
        updateHeight(node);

        // 检查是否失衡
        return balance(node);
    }



    /**
     * 检查节点是否失衡,失衡就进行旋转保持平衡
     * @param node 节点
     * @return 平衡后的节点
     */
    private AVLNode balance(AVLNode node){
        if (node == null){
            return null;
        }
        // 判断是否失衡,失衡的话判断是哪种失衡
        // 额外:还有一种情况,子树的高度等于0, 所以等于0也应该左右旋.
        int bf = bf(node);
        if (bf > 1 && bf(node.left) >= 0) { // LL
            return rightRotate(node);// 右旋保持平衡
        }else if (bf > 1 && bf(node.left) < 0){ // LR
            return leftRightRotate(node);// 左右旋保持平衡
        }else if (bf > 1 && bf(node.right) > 0){ // RL
            return rightLeftRotate(node);// 右左旋保持平衡
        }else if (bf > 1 && bf(node.right) <= 0){ // RR
            return leftRotate(node);// 左旋保持平衡
        }

        return node;
    }

    /**
     * 解决RR失衡情况
     * 右旋, 记得看图写代码
     * @param red 需要旋转的节点
     * @return 新的根节点
     */
    private AVLNode rightRotate(AVLNode red){
        // 得到需要旋转节点的左节点
        AVLNode yellow = red.left;
        // 需要换父节点的节点
        AVLNode green = yellow.right;
        // 进行右旋:黄色节点上位
        yellow.right = red;
        // 换父节点:绿色节点进行换爹
        red.left = green;
        // 旋转后更新节点高度
        updateHeight(red);
        updateHeight(yellow);
        return yellow;
    }

    /**
     * 解决LL失衡情况
     * 左旋, 记得看图写代码
     * @param red 需要旋转的节点
     * @return 新的根节点
     */
    private AVLNode leftRotate(AVLNode red){
        // 得到需要旋转节点的左节点
        AVLNode yellow = red.right;
        // 需要换父节点的节点
        AVLNode green = yellow.left;
        // 进行右旋:黄色节点上位
        yellow.left = red;
        // 换父节点:绿色节点进行换爹
        red.right = green;
        // 旋转后更新节点高度
        updateHeight(red);
        updateHeight(yellow);
        return yellow;
    }

    /**
     * 解决LR失衡情况
     * 左右旋, 先左旋左子树,再右根节点, 记得看图写代码
     * @param red 需要旋转的节点
     * @return 新的根节点
     */
    private AVLNode leftRightRotate(AVLNode red){
        // 得到失衡节点的左子树, 然后进行左旋
        red.left = leftRotate(red.left);
        // 根节点进行右旋
        return rightRotate(red);
    }

    /**
     * 解决RL失衡情况
     * 右左旋, 先右旋右子树, 再左根节点, 记得看图写代码
     * @param red 需要旋转的节点
     * @return 新的根节点
     */
    private AVLNode rightLeftRotate(AVLNode red){
        // 得到失衡节点的右子树, 然后进行右旋
        red.right = rightRotate(red.right);
        // 根节点进行左旋
        return leftRotate(red);
    }

    /**
     * 求节点高度
     * @param node 节点
     * @return 节点高度
     */
    private int height(AVLNode node) {
        return node == null ? 0 : node.height;
    }

    /**
     * 更新节点高度
     * @param node 需要被更新的节点
     */
    private void updateHeight(AVLNode node) {
        // 等于左右子树高度中大的一个再加一
        node.height = Integer.max(height(node.left), height(node.right)) + 1;
    }

    /**
     * 判断节点是否平衡的方法
     * @param node 需要判断的节点
     * @return 0，1，-1 时，表示左右平衡,  bf > 1 时，表示左边太高,  bf < -1 时，表示右边太高
     */
    private int bf(AVLNode node) {
        // 左节点高度 - 右节点高度
        return height(node.left) - height(node.right);
    }

}
